import { useCallback, useRef, useState } from 'react'
import { chatAPI } from '../services/api'

export interface ChatMessage {
  role: 'system' | 'user' | 'assistant'
  content: string
}

export function useSSEChat(initialMessages: ChatMessage[] = []) {
  const [messages, setMessages] = useState<ChatMessage[]>(initialMessages)
  const [isStreaming, setIsStreaming] = useState(false)
  const [error, setError] = useState<string | null>(null)
  const abortRef = useRef<AbortController | null>(null)

  const send = useCallback(async (
    input: string,
    options?: { sessionId?: string; context?: any }
  ) => {
    setError(null)
    abortRef.current?.abort()
    const abort = new AbortController()
    abortRef.current = abort

    // 确保消息类型正确
    const userMessage: ChatMessage = { role: 'user', content: input }
    const nextMessages: ChatMessage[] = [...messages, userMessage]
    setMessages(nextMessages)
    setIsStreaming(true)

    try {
      // 创建符合API要求的消息数组
      const apiMessages = nextMessages.map(msg => ({
        role: msg.role,
        content: msg.content
      }))
      
      const stream = await chatAPI.streamChat(apiMessages, { ...options, signal: abort.signal })
      const reader = stream.getReader()
      const decoder = new TextDecoder('utf-8')

      let assistantBuffer = ''
      setMessages(prev => [...prev, { role: 'assistant', content: '' }])

      const pushDelta = (delta: string) => {
        assistantBuffer += delta
        setMessages(prev => {
          const copy = [...prev]
          // 更新最后一条assistant消息
          for (let i = copy.length - 1; i >= 0; i--) {
            if (copy[i].role === 'assistant') {
              copy[i] = { ...copy[i], content: assistantBuffer }
              break
            }
          }
          return copy
        })
      }

      while (true) {
        const { value, done } = await reader.read()
        if (done) break
        const chunk = decoder.decode(value, { stream: true })

        // 解析SSE: 逐行读取 data: ...
        const lines = chunk.split(/\n\n/)
        for (const block of lines) {
          const line = block.trim()
          if (!line) continue
          if (line.startsWith('data:')) {
            const payload = line.slice(5).trim()
            if (payload === '[DONE]') {
              setIsStreaming(false)
              return
            }
            try {
              const data = JSON.parse(payload)
              if (data?.delta) pushDelta(data.delta)
              if (data?.final) setIsStreaming(false)
              if (data?.error) {
                setError(data.message || 'STREAM_ERROR')
                setIsStreaming(false)
              }
            } catch {
              // 非JSON数据，忽略
            }
          }
        }
      }
    } catch (e: any) {
      if (e?.name === 'AbortError') return
      setError(e?.message || 'STREAM_FAILED')
    } finally {
      setIsStreaming(false)
    }
  }, [messages])

  const cancel = useCallback(() => {
    abortRef.current?.abort()
    setIsStreaming(false)
  }, [])

  return { messages, setMessages, isStreaming, error, send, cancel }
}


